//! Physical memory frame management in Zig
//! SPDX-License-Identifier: MPL-2.0

const std = @import("std");
const lib = @import("lib.zig");

const Paddr = lib.Paddr;
const PAGE_SIZE = lib.PAGE_SIZE;
const MemoryError = lib.MemoryError;

// Frame metadata structure
pub const FrameMeta = struct {
    ref_count: std.atomic.Value(u32),
    paddr: Paddr,
    level: lib.PagingLevel,
    flags: FrameFlags,

    const Self = @This();

    pub fn init(paddr: Paddr, level: lib.PagingLevel) Self {
        return Self{
            .ref_count = std.atomic.Value(u32).init(1),
            .paddr = paddr,
            .level = level,
            .flags = FrameFlags{},
        };
    }

    pub fn incRef(self: *Self) void {
        _ = self.ref_count.fetchAdd(1, .acquire);
    }

    pub fn decRef(self: *Self) bool {
        const old_count = self.ref_count.fetchSub(1, .release);
        return old_count == 1;
    }

    pub fn refCount(self: *const Self) u32 {
        return self.ref_count.load(.acquire);
    }
};

pub const FrameFlags = packed struct {
    zeroed: bool = false,
    dma_coherent: bool = false,
    locked: bool = false,
    _reserved: u29 = 0,
};

// Frame handle
pub const Frame = struct {
    meta: *FrameMeta,

    const Self = @This();

    pub fn init(paddr: Paddr, frame_level: lib.PagingLevel, allocator: std.mem.Allocator) !Self {
        const meta = try allocator.create(FrameMeta);
        meta.* = FrameMeta.init(paddr, frame_level);
        return Self{ .meta = meta };
    }

    pub fn deinit(self: Self, allocator: std.mem.Allocator) void {
        if (self.meta.decRef()) {
            allocator.destroy(self.meta);
        }
    }

    pub fn clone(self: Self) Self {
        self.meta.incRef();
        return Self{ .meta = self.meta };
    }

    pub fn startPaddr(self: Self) Paddr {
        return self.meta.paddr;
    }

    pub fn endPaddr(self: Self) Paddr {
        return self.meta.paddr + self.size();
    }

    pub fn size(self: Self) usize {
        return lib.pageSize(self.meta.level);
    }

    pub fn level(self: Self) lib.PagingLevel {
        return self.meta.level;
    }

    pub fn asSlice(self: Self) []u8 {
        const vaddr = lib.paddrToVaddr(self.meta.paddr);
        return @as([*]u8, @ptrFromInt(vaddr))[0..self.size()];
    }

    pub fn asMutSlice(self: Self) []u8 {
        const vaddr = lib.paddrToVaddr(self.meta.paddr);
        return @as([*]u8, @ptrFromInt(vaddr))[0..self.size()];
    }

    pub fn zero(self: Self) void {
        const slice = self.asMutSlice();
        @memset(slice, 0);
        self.meta.flags.zeroed = true;
    }

    pub fn isZeroed(self: Self) bool {
        return self.meta.flags.zeroed;
    }
};

// Frame segment for contiguous frames
pub const Segment = struct {
    frames: []Frame,
    allocator: std.mem.Allocator,

    const Self = @This();

    pub fn init(start_paddr: Paddr, num_frames: usize, frame_level: lib.PagingLevel, allocator: std.mem.Allocator) !Self {
        const frames = try allocator.alloc(Frame, num_frames);
        errdefer allocator.free(frames);

        const frame_size = lib.pageSize(frame_level);
        for (frames, 0..) |*frame, i| {
            const paddr = start_paddr + i * frame_size;
            frame.* = try Frame.init(paddr, frame_level, allocator);
        }

        return Self{
            .frames = frames,
            .allocator = allocator,
        };
    }

    pub fn deinit(self: Self) void {
        for (self.frames) |frame| {
            frame.deinit(self.allocator);
        }
        self.allocator.free(self.frames);
    }

    pub fn startPaddr(self: Self) Paddr {
        if (self.frames.len == 0) return 0;
        return self.frames[0].startPaddr();
    }

    pub fn endPaddr(self: Self) Paddr {
        if (self.frames.len == 0) return 0;
        const last_frame = self.frames[self.frames.len - 1];
        return last_frame.endPaddr();
    }

    pub fn totalSize(self: Self) usize {
        if (self.frames.len == 0) return 0;
        return self.frames.len * self.frames[0].size();
    }

    pub fn nframes(self: Self) usize {
        return self.frames.len;
    }

    pub fn getFrame(self: Self, index: usize) ?Frame {
        if (index >= self.frames.len) return null;
        return self.frames[index].clone();
    }

    pub fn zero(self: Self) void {
        for (self.frames) |frame| {
            frame.zero();
        }
    }
};

// Frame allocation options
pub const FrameAllocOptions = struct {
    zeroed: bool = true,
    alignment: usize = PAGE_SIZE,
    level: lib.PagingLevel = 1,

    pub fn new() FrameAllocOptions {
        return FrameAllocOptions{};
    }

    pub fn withZeroed(self: FrameAllocOptions, zeroed: bool) FrameAllocOptions {
        var opts = self;
        opts.zeroed = zeroed;
        return opts;
    }

    pub fn withAlignment(self: FrameAllocOptions, alignment: usize) FrameAllocOptions {
        var opts = self;
        opts.alignment = alignment;
        return opts;
    }

    pub fn withLevel(self: FrameAllocOptions, level: lib.PagingLevel) FrameAllocOptions {
        var opts = self;
        opts.level = level;
        return opts;
    }
};

// Global frame allocator interface
pub const FrameAllocator = struct {
    allocFn: *const fn (self: *anyopaque, size: usize, alignment: usize) ?Paddr,
    deallocFn: *const fn (self: *anyopaque, paddr: Paddr, size: usize) void,
    ptr: *anyopaque,

    const Self = @This();

    pub fn alloc(self: Self, size: usize, alignment: usize) ?Paddr {
        return self.allocFn(self.ptr, size, alignment);
    }

    pub fn dealloc(self: Self, paddr: Paddr, size: usize) void {
        self.deallocFn(self.ptr, paddr, size);
    }

    pub fn allocFrame(self: Self, options: FrameAllocOptions, allocator: std.mem.Allocator) !Frame {
        const frame_size = lib.pageSize(options.level);
        const paddr = self.alloc(frame_size, options.alignment) orelse return MemoryError.OutOfMemory;

        var frame = try Frame.init(paddr, options.level, allocator);
        if (options.zeroed) {
            frame.zero();
        }
        return frame;
    }

    pub fn allocSegment(self: Self, nframes: usize, options: FrameAllocOptions, allocator: std.mem.Allocator) !Segment {
        const frame_size = lib.pageSize(options.level);
        const total_size = nframes * frame_size;
        const paddr = self.alloc(total_size, options.alignment) orelse return MemoryError.OutOfMemory;

        var segment = try Segment.init(paddr, nframes, options.level, allocator);
        if (options.zeroed) {
            segment.zero();
        }
        return segment;
    }

    pub fn deallocFrame(self: Self, frame: Frame) void {
        const size = frame.size();
        const paddr = frame.startPaddr();
        self.dealloc(paddr, size);
    }

    pub fn deallocSegment(self: Self, segment: Segment) void {
        const size = segment.totalSize();
        const paddr = segment.startPaddr();
        self.dealloc(paddr, size);
    }
};

// Simple buddy allocator implementation
pub const BuddyAllocator = struct {
    memory_start: Paddr,
    memory_size: usize,
    free_lists: [32]?*FreeBlock, // Support up to 2^31 bytes
    mutex: std.Thread.Mutex,

    const FreeBlock = struct {
        next: ?*FreeBlock,
        size: usize,
    };

    const Self = @This();

    pub fn init(memory_start: Paddr, memory_size: usize) Self {
        var allocator = Self{
            .memory_start = memory_start,
            .memory_size = memory_size,
            .free_lists = [_]?*FreeBlock{null} ** 32,
            .mutex = std.Thread.Mutex{},
        };

        // Initialize with one large block
        const order = std.math.log2(memory_size);
        const block = @as(*FreeBlock, @ptrFromInt(lib.paddrToVaddr(memory_start)));
        block.* = FreeBlock{ .next = null, .size = memory_size };
        allocator.free_lists[order] = block;

        return allocator;
    }

    pub fn frameAllocator(self: *Self) FrameAllocator {
        return FrameAllocator{
            .allocFn = allocImpl,
            .deallocFn = deallocImpl,
            .ptr = self,
        };
    }

    fn allocImpl(ptr: *anyopaque, size: usize, alignment: usize) ?Paddr {
        const self: *Self = @ptrCast(@alignCast(ptr));
        return self.allocAligned(size, alignment);
    }

    fn deallocImpl(ptr: *anyopaque, paddr: Paddr, size: usize) void {
        const self: *Self = @ptrCast(@alignCast(ptr));
        self.deallocAligned(paddr, size);
    }

    fn allocAligned(self: *Self, size: usize, alignment: usize) ?Paddr {
        self.mutex.lock();
        defer self.mutex.unlock();

        const aligned_size = lib.alignUp(size, alignment);
        const order = std.math.log2(aligned_size);

        // Find a suitable block
        for (order..self.free_lists.len) |i| {
            if (self.free_lists[i]) |block| {
                // Remove from free list
                self.free_lists[i] = block.next;

                // Split if necessary
                var current_size = @as(usize, 1) << @intCast(i);
                var current_addr = lib.vaddrToPaddr(@intFromPtr(block));

                while (current_size > aligned_size) {
                    current_size /= 2;
                    const buddy_addr = current_addr + current_size;
                    const buddy_block = @as(*FreeBlock, @ptrFromInt(lib.paddrToVaddr(buddy_addr)));
                    buddy_block.* = FreeBlock{ .next = self.free_lists[std.math.log2(current_size)], .size = current_size };
                    self.free_lists[std.math.log2(current_size)] = buddy_block;
                }

                return current_addr;
            }
        }

        return null;
    }

    fn deallocAligned(self: *Self, paddr: Paddr, size: usize) void {
        self.mutex.lock();
        defer self.mutex.unlock();

        var current_addr = paddr;
        var current_size = size;
        var order = std.math.log2(current_size);

        // Try to merge with buddy
        while (order < self.free_lists.len - 1) {
            const buddy_addr = current_addr ^ current_size;
            
            // Check if buddy is free
            var prev: ?*FreeBlock = null;
            var current = self.free_lists[order];
            
            while (current) |block| {
                if (lib.vaddrToPaddr(@intFromPtr(block)) == buddy_addr and block.size == current_size) {
                    // Remove buddy from free list
                    if (prev) |p| {
                        p.next = block.next;
                    } else {
                        self.free_lists[order] = block.next;
                    }

                    // Merge
                    current_addr = @min(current_addr, buddy_addr);
                    current_size *= 2;
                    order += 1;
                    break;
                }
                prev = block;
                current = block.next;
            } else {
                break;
            }
        }

        // Add merged block to free list
        const block = @as(*FreeBlock, @ptrFromInt(lib.paddrToVaddr(current_addr)));
        block.* = FreeBlock{ .next = self.free_lists[order], .size = current_size };
        self.free_lists[order] = block;
    }
};

// Tests
test "frame creation and reference counting" {
    const testing = std.testing;
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    const allocator = gpa.allocator();

    var frame = try Frame.init(0x1000, 1, allocator);
    defer frame.deinit(allocator);

    try testing.expectEqual(@as(Paddr, 0x1000), frame.startPaddr());
    try testing.expectEqual(@as(Paddr, 0x2000), frame.endPaddr());
    try testing.expectEqual(PAGE_SIZE, frame.size());

    const cloned_frame = frame.clone();
    defer cloned_frame.deinit(allocator);

    try testing.expectEqual(@as(u32, 2), frame.meta.refCount());
}

test "segment creation" {
    const testing = std.testing;
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    const allocator = gpa.allocator();

    var segment = try Segment.init(0x10000, 4, 1, allocator);
    segment.deinit();

    try testing.expectEqual(@as(usize, 4), segment.nframes());
    try testing.expectEqual(@as(Paddr, 0x10000), segment.startPaddr());
    try testing.expectEqual(@as(Paddr, 0x14000), segment.endPaddr());
    try testing.expectEqual(@as(usize, 4 * PAGE_SIZE), segment.totalSize());
}